package com.ocean.everyday.d002_findingmkaverage;

import java.util.Deque;
import java.util.LinkedList;
import java.util.TreeMap;

/**
 * https://leetcode.cn/problems/finding-mk-average/
 * 求出MK平均值
 */
public class MKAverage3 {
    private TreeMap<Integer, Integer> treeMap1; // 存放最小的k个值
    private TreeMap<Integer, Integer> treeMap2; // 存放最小最大之外的值
    private TreeMap<Integer, Integer> treeMap3; // 存放最大的k个值
    private int size1; // 最小值的个数
    private int size2; // 中间元素个数
    private int size3; // 最大值的个数
    private int sum; // 剔除最大值 最小值后生效元素的和
    private Deque<Integer> deque; // 存储所有的元素队列
    private int m;
    private int k;

    public static void main(String[] args) {
        MKAverage3 obj = new MKAverage3(3, 1);
        obj.addElement(3);        // 当前元素为 [3]
        obj.addElement(1);        // 当前元素为 [3,1]
        int i = obj.calculateMKAverage();// 返回 -1 ，因为 m = 3
        System.out.println(i);
        obj.addElement(10);
        i = obj.calculateMKAverage();
        System.out.println(i);

        obj.addElement(5);        // 当前元素为 [3,1,10,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5]
        obj.addElement(5);        // 当前元素为 [3,1,10,5,5,5]
        i = obj.calculateMKAverage();
        System.out.println(i);
    }


    public MKAverage3(int m, int k) {
        this.treeMap1 = new TreeMap<>();
        this.treeMap2 = new TreeMap<>();
        this.treeMap3 = new TreeMap<>();
        this.size1 = 0;
        this.size2 = 0;
        this.size3 = 0;
        this.deque = new LinkedList<>();
        this.sum = 0;
        this.m = m;
        this.k = k;
    }

    public void addElement(int num) {
        deque.offer(num);
        // 判断元素个数
        if (deque.size() <= this.m) {
            put(treeMap2, num);
            // 求和
            size2++;
            sum += num;
            if (deque.size() == this.m) {
                // 要将元素移动到最小最大集合中
                while (size1 < k) {
                    Integer firstKey = treeMap2.firstKey();
                    put(treeMap1, firstKey);
                    size1++;
                    remove(treeMap2, firstKey);
                    size2--;
                    sum -= firstKey;

                }
                while (size3 < k) {
                    Integer lastKey = treeMap2.lastKey();
                    put(treeMap3, lastKey);
                    size3++;
                    remove(treeMap2, lastKey);
                    size2--;
                    sum -= lastKey;
                }
            }
            return;
        }
        // 先把当前元素加入到集合中
        if (num < treeMap1.lastKey()) {
            Integer last = treeMap1.lastKey();
            put(treeMap1, num);
            remove(treeMap1, last);
            put(treeMap2, last);
            size2++;
            sum += last;
        } else if (num > treeMap3.firstKey()) {
            Integer first = treeMap3.firstKey();
            put(treeMap3, num);
            remove(treeMap3, first);
            put(treeMap2, first);
            size2++;
            sum += first;
        } else {
            put(treeMap2, num);
            size2++;
            sum += num;
        }
        // 剔除一个
        Integer poll = deque.poll();

        if (treeMap1.containsKey(poll)) {
            Integer firstKey = treeMap2.firstKey();
            remove(treeMap1, poll);
            put(treeMap1, firstKey);
            remove(treeMap2, firstKey);
            size2--;
            sum -= firstKey;
        } else if (treeMap3.containsKey(poll)) {
            Integer lastKey = treeMap2.lastKey();
            remove(treeMap3, poll);
            put(treeMap3, lastKey);
            remove(treeMap2, lastKey);
            size2--;
            sum -= lastKey;
        } else {
            remove(treeMap2, poll);
            size2--;
            sum -= poll;
        }

    }

    /**
     * @return
     */
    public int calculateMKAverage() {
        if (deque.size() < m) {
            return -1;
        }
        return (int) sum / (m - 2 * k);
    }

    private void put(TreeMap<Integer, Integer> treeMap, int num) {
        treeMap.put(num, treeMap.getOrDefault(num, 0) + 1);
    }

    private void remove(TreeMap<Integer, Integer> treeMap, int num) {
        treeMap.put(num, treeMap.get(num) - 1);
        if (treeMap.get(num) == 0) {
            treeMap.remove(num);
        }
    }
}
